package ssau.knyazev.ocean;

import static jcuda.driver.CUgraphicsMapResourceFlags.CU_GRAPHICS_MAP_RESOURCE_FLAGS_WRITE_DISCARD;
import static jcuda.driver.JCudaDriver.*;

import java.awt.BorderLayout;
import java.awt.Dimension;
import java.awt.Point;
import java.awt.event.*;
import java.io.*;
import java.nio.*;
import java.util.Random;

import javax.media.opengl.*;
import javax.media.opengl.awt.GLCanvas;
import javax.media.opengl.glu.GLU;
import javax.swing.JFrame;
import javax.swing.SwingUtilities;

import ssau.knyazev.ocean.OceanConst;

import jcuda.Pointer;
import jcuda.Sizeof;
import jcuda.driver.*;
import jcuda.jcufft.JCufft;
import jcuda.jcufft.cufftHandle;
import jcuda.jcufft.cufftType;


import com.jogamp.opengl.util.Animator;

public class Ocean implements GLEventListener, OceanConst {
	
	private JFrame frame = null;
	private Animator animator = null;
	private GL2 gl = null;
	
	//OpenGL variables
	private int shaderProgramID = 0;
	
	private float translateX = 0.0f;
	private float translateY = 0.0f;
	private float translateZ = -2.0f;
	private float rotateX = 20f;
	private float rotateY = 0.0f;
	
	//OpenGL & Cuda variables
	private int posVertexBuffer = 0;	
	private int heightVertexBuffer = 0;
	private int slopeVertexBuffer = 0;	
	private int indexBuffer = 0;
	
	private CUgraphicsResource cuda_heightVB_resource = null;	
	private CUgraphicsResource cuda_slopeVB_resource = null;
	
	// FFT data
	private cufftHandle fftPlan = null;
	private CUdeviceptr d_h0Ptr = null;
	private CUdeviceptr d_htPtr = null;
	private CUdeviceptr d_slopePtr = null;
	private float[] h_h0 = null;	
	
	//JCUDA
	private CUdevice device = null;
	private CUcontext glContext = null;
	private CUfunction generateSpectrumKernel = null;
	private CUfunction updateHeightmapKernel = null;
	private CUfunction calculateSlopeKernel = null;
	private CUmodule module = null;
		
	private float animTime = 0;
	
	private static long sum = 0;
	private static int cur = 1;
	private static int count = 50;
		
	
	public Ocean(GLCapabilities capabilities) {
		GLCanvas glComponent = new GLCanvas(capabilities);
		glComponent.setFocusable(true);
		glComponent.addGLEventListener(this);

		MouseControl mouseControl = new MouseControl();
		glComponent.addMouseMotionListener(mouseControl);
		glComponent.addMouseWheelListener(mouseControl);
		
		frame = new JFrame("WaterSurface");
		frame.addWindowListener(new WindowAdapter() {
			@Override
			public void windowClosing(WindowEvent e) {
				animator.stop();
				try{
					cuMemFree(d_h0Ptr);
					cuMemFree(d_htPtr);
					cuMemFree(d_slopePtr);
					JCufft.cufftDestroy(fftPlan);
					gl.glDeleteBuffers(1, IntBuffer.wrap(new int[] {posVertexBuffer}));
					gl.glDeleteBuffers(1, IntBuffer.wrap(new int[] {heightVertexBuffer}));
					gl.glDeleteBuffers(1, IntBuffer.wrap(new int[] {slopeVertexBuffer}));

				} catch (Exception e1){
					e1.printStackTrace();
				}
				System.exit(0);
			}
		});
		frame.setLayout(new BorderLayout());
		glComponent.setPreferredSize(new Dimension(800, 800));
		frame.add(glComponent, BorderLayout.CENTER);
		frame.pack();
		frame.setVisible(true);
		glComponent.requestFocus();

		// Create and start the animator
		animator = new Animator(glComponent);
		animator.start();

	}

	public static void main(String args[]) {
		GLProfile profile = GLProfile.get(GLProfile.GL2);
		final GLCapabilities capabilities = new GLCapabilities(profile);
		SwingUtilities.invokeLater(new Runnable() {
			public void run() {
				new Ocean(capabilities);
			}
		});
	}


	@Override
	public void init(GLAutoDrawable drawable) {
		gl = drawable.getGL().getGL2();
		gl.setSwapInterval(0);
		gl.glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
		gl.glEnable(GL2.GL_DEPTH_TEST);

		initShaders(gl);
		initJCuda();
		initVBO(gl);
	}
	
	private void initJCuda(){
		setExceptionsEnabled(true);
		cuInit(0);
		device = new CUdevice();
		cuDeviceGet(device, 0);
		glContext = new CUcontext();
		cuGLCtxCreate(glContext, CUctx_flags.CU_CTX_BLOCKING_SYNC, device);

		String ptxFileName = "";
		try {
			ptxFileName = preparePtxFile("resources/ocean.cu");
		} catch (IOException e) {
			System.err.println("Could not create PTX file");
			throw new RuntimeException("Could not create PTX file", e);
		}

		module = new CUmodule();
		cuModuleLoad(module, ptxFileName);
		generateSpectrumKernel = new CUfunction();
		cuModuleGetFunction(generateSpectrumKernel, module, "generateSpectrumKernel");
		updateHeightmapKernel = new CUfunction();
		cuModuleGetFunction(updateHeightmapKernel, module, "updateHeightmapKernel");
		calculateSlopeKernel = new CUfunction();
		cuModuleGetFunction(calculateSlopeKernel, module, "calculateSlopeKernel");

		fftPlan = new cufftHandle();
		JCufft.cufftPlan2d(fftPlan, meshSize, meshSize, cufftType.CUFFT_C2C);

		int spectrumSize = spectrumW * spectrumH * 2;
		h_h0 = new float[spectrumSize];
		h_h0 = generate_h0(h_h0);
		d_h0Ptr = new CUdeviceptr();
		cuMemAlloc(d_h0Ptr, h_h0.length*Sizeof.FLOAT);
		cuMemcpyHtoD(d_h0Ptr, Pointer.to(h_h0), h_h0.length*Sizeof.FLOAT);
		
		int outputSize =  meshSize*meshSize*Sizeof.FLOAT*2;
		d_htPtr = new CUdeviceptr();
		d_slopePtr = new CUdeviceptr();
		cuMemAlloc(d_htPtr, outputSize);
		cuMemAlloc(d_slopePtr, outputSize);
	}

	private void initVBO(GL2 gl) {		
		int[] buffer = new int[1];
		int size = meshSize*meshSize*Sizeof.FLOAT;
		gl.glGenBuffers(1, IntBuffer.wrap(buffer));
		heightVertexBuffer = buffer[0];
		gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, heightVertexBuffer);
		gl.glBufferData(GL2.GL_ARRAY_BUFFER, size, (Buffer) null, GL2.GL_DYNAMIC_DRAW);
		gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, 0);
		cuda_heightVB_resource = new CUgraphicsResource();
		cuGraphicsGLRegisterBuffer(cuda_heightVB_resource, heightVertexBuffer, CU_GRAPHICS_MAP_RESOURCE_FLAGS_WRITE_DISCARD);
		
		buffer = new int[1];
		size = meshSize*meshSize*Sizeof.FLOAT*2;
		gl.glGenBuffers(1, IntBuffer.wrap(buffer));
		slopeVertexBuffer = buffer[0];
		gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, slopeVertexBuffer);
		gl.glBufferData(GL2.GL_ARRAY_BUFFER, size, (Buffer) null, GL2.GL_DYNAMIC_DRAW);
		gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, 0);
		cuda_slopeVB_resource = new CUgraphicsResource();
		cuGraphicsGLRegisterBuffer(cuda_slopeVB_resource, slopeVertexBuffer, CU_GRAPHICS_MAP_RESOURCE_FLAGS_WRITE_DISCARD);

		buffer = new int[1];
		size = meshSize*meshSize*Sizeof.FLOAT*4;
		gl.glGenBuffers(1, IntBuffer.wrap(buffer));
		posVertexBuffer = buffer[0];
		gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, posVertexBuffer);
		gl.glBufferData(GL2.GL_ARRAY_BUFFER, size, (Buffer) null, GL2.GL_DYNAMIC_DRAW);
		gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, 0);
		gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, posVertexBuffer);
		ByteBuffer byteBuffer = gl.glMapBuffer(GL2.GL_ARRAY_BUFFER, GL2.GL_WRITE_ONLY);
		if (byteBuffer != null){
			FloatBuffer put = byteBuffer.asFloatBuffer();
			int index = 0;
			for (int y = 0; y < meshSize; y++) {
				for (int x = 0; x < meshSize; x++) {
					float u = x / (float) (meshSize - 1);
					float v = y / (float) (meshSize - 1);
					put.put(index, u * 2.0f - 1.0f);
					put.put(index+1, 0.0f);
					put.put(index + 2, v * 2.0f - 1.0f);
					put.put(index+3, 1.0f);
					index += 4;
				}
			}
		}
		gl.glUnmapBuffer(GL2.GL_ARRAY_BUFFER);
		gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, 0);

		size = ((meshSize*2)+2)*(meshSize-1)*Sizeof.INT;
		buffer = new int[1];
		gl.glGenBuffers(1, IntBuffer.wrap(buffer));
		indexBuffer = buffer[0];
		gl.glBindBuffer(GL2.GL_ELEMENT_ARRAY_BUFFER, indexBuffer);
		gl.glBufferData(GL2.GL_ELEMENT_ARRAY_BUFFER, size, (Buffer) null, GL2.GL_STATIC_DRAW);
		
		byteBuffer = gl.glMapBuffer(GL2.GL_ELEMENT_ARRAY_BUFFER, GL2.GL_WRITE_ONLY);
		if (byteBuffer != null){
			IntBuffer indices = byteBuffer.asIntBuffer();
			int index = 0;
			for(int y=0; y<meshSize-1; y++) {
		        for(int x=0; x<meshSize; x++) {
		        	indices.put(index, y*meshSize+x);
		        	indices.put(index+1, (y+1)*meshSize+x);
		        	index +=2;
		        }
		        indices.put(index, (y+1)*meshSize+(meshSize-1));
		        indices.put(index+1, (y+1)*meshSize);
		        index += 2;
		    }
		}
		gl.glUnmapBuffer(GL2.GL_ELEMENT_ARRAY_BUFFER);
		gl.glBindBuffer(GL2.GL_ELEMENT_ARRAY_BUFFER, 0);
	}

	@Override
	public void display(GLAutoDrawable drawable) {
		float delta = System.nanoTime();
		long time = System.nanoTime();
		gl = drawable.getGL().getGL2();

		runCuda();

		gl.glClear(GL2.GL_COLOR_BUFFER_BIT | GL2.GL_DEPTH_BUFFER_BIT);
		
		gl.glMatrixMode(GL2.GL_MODELVIEW);
		gl.glLoadIdentity();
		gl.glTranslatef(translateX, translateY, translateZ);
		gl.glRotatef(rotateX, 1.0f, 0.0f, 0.0f);
		gl.glRotatef(rotateY, 0.0f, 1.0f, 0.0f);

		gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, posVertexBuffer);
		gl.glVertexPointer(4, GL2.GL_FLOAT, 0, 0);
		gl.glEnableClientState(GL2.GL_VERTEX_ARRAY);

		gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, heightVertexBuffer);
		gl.glClientActiveTexture(GL2.GL_TEXTURE0);
		gl.glTexCoordPointer(1, GL2.GL_FLOAT, 0, 0);
		gl.glEnableClientState(GL2.GL_TEXTURE_COORD_ARRAY);

		gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, slopeVertexBuffer);
		gl.glClientActiveTexture(GL2.GL_TEXTURE1);
		gl.glTexCoordPointer(2, GL2.GL_FLOAT, 0, 0);
		gl.glEnableClientState(GL2.GL_TEXTURE_COORD_ARRAY);

		gl.glUseProgram(shaderProgramID);

		int uniHeightScale = gl.glGetUniformLocation(shaderProgramID, "heightScale");
		gl.glUniform1f(uniHeightScale, 0.5f);

		int uniChopiness = gl.glGetUniformLocation(shaderProgramID, "chopiness");
		gl.glUniform1f(uniChopiness, 1.0f);

		int uniSize = gl.glGetUniformLocation(shaderProgramID, "size");
		gl.glUniform2f(uniSize, (float) meshSize, (float) meshSize);

		int uniDeepColor = gl.glGetUniformLocation(shaderProgramID, "deepColor");
		gl.glUniform4f(uniDeepColor, 0.0f, 0.1f, 0.4f, 1.0f);

		int uniShallowColor = gl.glGetUniformLocation(shaderProgramID, "shallowColor");
		gl.glUniform4f(uniShallowColor, 0.1f, 0.3f, 0.3f, 1.0f);

		int uniSkyColor = gl.glGetUniformLocation(shaderProgramID, "skyColor");
		gl.glUniform4f(uniSkyColor, 1.0f, 1.0f, 1.0f, 1.0f);

		int uniLightDir = gl.glGetUniformLocation(shaderProgramID, "lightDir");
		gl.glUniform3f(uniLightDir, 0.0f, 1.0f, 0.0f);
		
		gl.glColor3f(1.0f, 1.0f, 1.0f);
		gl.glBindBuffer(GL2.GL_ELEMENT_ARRAY_BUFFER, indexBuffer);
		gl.glPolygonMode(GL2.GL_FRONT_AND_BACK, GL2.GL_FILL);
		gl.glDrawElements(GL2.GL_TRIANGLE_STRIP, ((meshSize * 2) + 2) * (meshSize - 1), GL2.GL_UNSIGNED_INT, 0);
		gl.glPolygonMode(GL2.GL_FRONT_AND_BACK, GL2.GL_FILL);
		gl.glBindBuffer(GL2.GL_ELEMENT_ARRAY_BUFFER, 0);
		
		gl.glDisableClientState(GL2.GL_VERTEX_ARRAY);
		gl.glClientActiveTexture(GL2.GL_TEXTURE0);
		gl.glDisableClientState(GL2.GL_TEXTURE_COORD_ARRAY);
		gl.glClientActiveTexture(GL2.GL_TEXTURE1);
		gl.glDisableClientState(GL2.GL_TEXTURE_COORD_ARRAY);

		gl.glUseProgram(0);


		drawable.swapBuffers();
		time = System.nanoTime() - time;
		sum +=time;
		cur++;
		if (cur == count){
			time = sum/count;
			System.out.println("DISPLAY TIME = " + time);
		}
		
		delta = System.nanoTime() - delta;
		animTime +=  delta/1000000000;

	}

	@Override
	public void dispose(GLAutoDrawable drawable) { }
	
	@Override
	public void reshape(GLAutoDrawable drawable, int x, int y, int width, int height) {
		gl = drawable.getGL().getGL2();
		GLU glu = GLU.createGLU(gl);
		gl.glViewport(0, 0, width, height);

		gl.glMatrixMode(GL2.GL_PROJECTION);
		gl.glLoadIdentity();
		glu.gluPerspective(60.0, (double) width / (double) height, 0.1, 10.0);
	}
	
	private void runCuda() {
		Pointer kernelParameters = null;
		int blockX = 8;
		int blockY = 8;
		int gridX = meshSize/blockX;
		int gridY = gridX;
		kernelParameters = Pointer.to(
				Pointer.to(d_h0Ptr),
				Pointer.to(d_htPtr),
				Pointer.to(new int[] { spectrumW }),
				Pointer.to(new int[] { meshSize }),
				Pointer.to(new int[] { meshSize }),
				Pointer.to(new float[] { animTime }),
				Pointer.to(new float[] { patchSize }));

		
		cuLaunchKernel(generateSpectrumKernel, 
				gridX, gridY, 1, // Grid dimension
				blockX, blockY, 1, // Block dimension
				0, null, // Shared memory size and stream
				kernelParameters, null // Kernel- and extra parameters
		);
		cuCtxSynchronize();
		
		JCufft.cufftExecC2C(fftPlan, d_htPtr, d_htPtr, JCufft.CUFFT_INVERSE);		
		
		CUdeviceptr g_hptr = new CUdeviceptr();
		cuGraphicsMapResources(1, new CUgraphicsResource[]{cuda_heightVB_resource}, null);
		cuGraphicsResourceGetMappedPointer(g_hptr, new long[1], cuda_heightVB_resource);
		kernelParameters = Pointer.to(
				Pointer.to(g_hptr),
				Pointer.to(d_htPtr),
				Pointer.to(new int[] { meshSize }));

		cuLaunchKernel(updateHeightmapKernel, 
				gridX, gridY, 1, // Grid dimension
				blockX, blockY, 1, // Block dimension
				0, null, // Shared memory size and stream
				kernelParameters, null // Kernel- and extra parameters
		);
		cuCtxSynchronize();
		cuGraphicsUnmapResources(1, new CUgraphicsResource[]{cuda_heightVB_resource}, null);
		
		CUdeviceptr g_sptr = new CUdeviceptr();
		cuGraphicsMapResources(1, new CUgraphicsResource[]{cuda_slopeVB_resource}, null);
		cuGraphicsResourceGetMappedPointer(g_sptr, new long[1], cuda_slopeVB_resource);
		kernelParameters = Pointer.to(
				Pointer.to(g_hptr),
				Pointer.to(g_sptr),
				Pointer.to(new int[] { meshSize }),
				Pointer.to(new int[] { meshSize }));

		cuLaunchKernel(calculateSlopeKernel, 
				gridX, gridY, 1, // Grid dimension
				blockX, blockY, 1, // Block dimension
				0, null, // Shared memory size and stream
				kernelParameters, null // Kernel- and extra parameters
		);
		cuCtxSynchronize();
		cuGraphicsUnmapResources(1, new CUgraphicsResource[]{ cuda_slopeVB_resource}, null);
	}
	
	private void initShaders(GL2 gl) {
		shaderProgramID = gl.glCreateProgram();
		attachShader(gl, GL2.GL_VERTEX_SHADER, vertexShaderSource);
		attachShader(gl, GL2.GL_FRAGMENT_SHADER, fragmentShaderSource);
		gl.glLinkProgram(shaderProgramID);
		
		int[] buffer = new int[1];
		gl.glGetProgramiv(shaderProgramID, GL2.GL_LINK_STATUS, IntBuffer.wrap(buffer));
		gl.glValidateProgram(shaderProgramID);
	}
	
	private int attachShader(GL2 gl, int type, String shaderSource){
		int shader = 0; 
		shader = gl.glCreateShader(type);
		gl.glShaderSource(shader, 1, new String[] { shaderSource }, null);
		gl.glCompileShader(shader);
		int[] buffer = new int[1];
		gl.glGetShaderiv(shader, GL2.GL_COMPILE_STATUS, IntBuffer.wrap(buffer));
		gl.glAttachShader(shaderProgramID, shader);
		gl.glDeleteShader(shader);
		return shader;
	}
	
	// Phillips spectrum
	// (Kx, Ky) - normalized wave vector
	// Vdir - wind angle in radians
	// V - wind speed
	// A - constant
	private float phillips(float Kx, float Ky, float Vdir, float V, float A, float dir_depend) {
		float k_squared = Kx*Kx + Ky*Ky;
		if (k_squared == 0.0f)
			return 0.0f;

		float L = V*V/g;

		float k_x = (float) (Kx / Math.sqrt(k_squared));
		float k_y = (float) (Ky / Math.sqrt(k_squared));
		float w_dot_k = (float) (k_x*Math.cos(Vdir) + k_y*Math.sin(Vdir));

		float phillips = (float) (A*Math.exp(-1/(k_squared*L*L))/(k_squared * k_squared) * w_dot_k*w_dot_k);

		// filter out waves moving opposite to wind
		if (w_dot_k < 0.0f)
			phillips *= dir_depend;

		return phillips;
	}

	private float[] generate_h0(float[] h0) {
		Random rnd = new Random();
		for (int y = 0; y <= meshSize; y++) {
			for (int x = 0; x <= meshSize; x++) {
				float kx = (float) ((-meshSize/2 + x) * (2*Math.PI / patchSize));
				float ky = (float) ((-meshSize/2 + y) * (2*Math.PI / patchSize));

				float P = (float) (Math.sqrt(phillips(kx, ky, windDir, windSpeed, A, dirDepend)));
				if (kx == 0.0f && ky == 0.0f){
					P = 0.0f;
				}
				float Er = (float) rnd.nextGaussian();
				float Ei = (float) rnd.nextGaussian();
				float h0_re = (float) (Er*P * Math.sqrt(2));
				float h0_im = (float) (Ei*P * Math.sqrt(2));

				int i = y * spectrumW + x;
				h0[2*i] = h0_re;
				h0[2*i + 1] = h0_im;
			}
		}
		return h0;
	}
	
	class MouseControl implements MouseMotionListener, MouseWheelListener {
		private Point previousMousePosition = new Point();

		@Override
		public void mouseDragged(MouseEvent e) {
			int dx = e.getX() - previousMousePosition.x;
			int dy = e.getY() - previousMousePosition.y;

			if ((e.getModifiersEx() & MouseEvent.BUTTON1_DOWN_MASK) == MouseEvent.BUTTON1_DOWN_MASK) {
				translateX += dx / 100.0f;
				translateY -= dy / 100.0f;
			} else if ((e.getModifiersEx() & MouseEvent.BUTTON3_DOWN_MASK) == MouseEvent.BUTTON3_DOWN_MASK) {
				rotateX += dy;
				rotateY += dx;
			}
			previousMousePosition = e.getPoint();

		}

		@Override
		public void mouseMoved(MouseEvent e) {
			previousMousePosition = e.getPoint();
		}

		@Override
		public void mouseWheelMoved(MouseWheelEvent e) {
			translateZ += e.getWheelRotation() * 0.25f;
			previousMousePosition = e.getPoint();
		}
	}
	
	private static String preparePtxFile(String cuFileName) throws IOException {
		int endIndex = cuFileName.lastIndexOf('.');
		if (endIndex == -1) {
			endIndex = cuFileName.length() - 1;
		}
		String ptxFileName = cuFileName.substring(0, endIndex + 1) + "ptx";

		File cuFile = new File(cuFileName);
		if (!cuFile.exists()) {
			throw new IOException("Input file not found: " + cuFileName);
		}
		String modelString = "-m" + System.getProperty("sun.arch.data.model");
		String command = "nvcc " + modelString + " -ptx " + cuFile.getPath() + " -o " + ptxFileName;

		System.out.println("Executing\n" + command);
		Process process = Runtime.getRuntime().exec(command);

		String errorMessage = new String(toByteArray(process.getErrorStream()));
		String outputMessage = new String(toByteArray(process.getInputStream()));
		int exitValue = 0;
		try {
			exitValue = process.waitFor();
		} catch (InterruptedException e) {
			Thread.currentThread().interrupt();
			throw new IOException("Interrupted while waiting for nvcc output", e);
		}

		if (exitValue != 0) {
			System.out.println("nvcc process exitValue " + exitValue);
			System.out.println("errorMessage:\n" + errorMessage);
			System.out.println("outputMessage:\n" + outputMessage);
			throw new IOException("Could not create .ptx file: " + errorMessage);
		}
		System.out.println("Finished creating PTX file");
		return ptxFileName;
	}

	private static byte[] toByteArray(InputStream inputStream) throws IOException {
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		byte buffer[] = new byte[8192];
		while (true) {
			int read = inputStream.read(buffer);
			if (read == -1) {
				break;
			}
			baos.write(buffer, 0, read);
		}
		return baos.toByteArray();
	}
	
}
